from __future__ import print_function, division
import os
import sys
import gzip
import random
import argparse

parser = argparse.ArgumentParser(description="Random pick n pair-end reads")
parser.add_argument("-i", "--infile1", help="input file1", required=True)
parser.add_argument("-I", "--infile2", help="input file2", required=True)
parser.add_argument("-o", "--outfile1", help="output file1", required=True)
parser.add_argument("-O", "--outfile2", help="output file2", required=True)
parser.add_argument("-p", "--pvalue", help="p value",
                    required=True, type=float)
args = vars(parser.parse_args())

pvalue = float(args['pvalue'])


def save_open(f, mode='r'):
    if f.endswith(".gz"):
        try:
            return gzip.open(f, mode)
        except:
            pass
    return open(f, mode)


if 'infile2' in args:
    print("PE reads")
    infile1 = save_open(args['infile1'])
    infile2 = save_open(args['infile2'])
    fo1 = save_open(args['outfile1'], 'w')
    fo2 = save_open(args['outfile2'], 'w')
    #reads = []
    #picked = {}
    count = 0
    idnum = 0
    for line in infile1:
        count += 1
        if count == 1:
            idline = line
            idline2 = infile2.readline()
        elif count == 2:
            seqline = line
            seqline2 = infile2.readline()
        elif count == 3:
            tline = line
            tline2 = infile2.readline()
        elif count == 4:
            qline = line
            qline2 = infile2.readline()
            count = 0
            if random.random() < pvalue:
                fo1.write(idline+seqline+tline+qline)
                fo2.write(idline2+seqline2+tline2+qline2)
